from data_access.database import Session
from sqlalchemy.exc import SQLAlchemyError

from data_access.models.models import Task
from service.base_service.logger_service import logger


class TaskDAO:
    def __init__(self):
        self.session = Session()

    def add_task(self, dispatch_code, task_id, task_name, task_frequency, task_status, sql_id, execute_database_id,
                 receive_service_id, storage_database_id, storage_table_id, comment):
        new_task = Task(
            dispatch_code=dispatch_code,
            task_id=task_id,
            task_name=task_name,
            task_frequency=task_frequency,
            task_status=task_status,
            sql_id=sql_id,
            execute_database_id=execute_database_id,
            receive_service_id=receive_service_id,
            storage_database_id=storage_database_id,
            storage_table_id=storage_table_id,
            comment=comment
        )
        try:
            self.session.add(new_task)
            self.session.commit()
            return True
        except SQLAlchemyError as e:
            self.session.rollback()
            return False
        finally:
            self.session.close()

    def get_task_by_task_id(self, dispatch_code, task_id):
        try:
            return self.session.query(Task).filter(Task.dispatch_code == dispatch_code, Task.task_id == task_id).first()
        except SQLAlchemyError as e:
            return None
        finally:
            self.session.close()

    def get_all_tasks(self):
        tasks = self.session.query(Task).all()
        return tasks

    def update_task_by_id(self, id, **kwargs):
        try:
            self.session.query(Task).filter(Task.id == id).update(kwargs)
            self.session.commit()
            return True
        except SQLAlchemyError as e:
            self.session.rollback()
            logger.error(f'Failed to update task: {e}')
            return False
        finally:
            self.session.close()

    def delete_task(self, task_id):
        try:
            self.session.query(Task).filter(Task.task_id == task_id).delete()
            self.session.commit()
            return True
        except SQLAlchemyError as e:
            self.session.rollback()
            return False
        finally:
            self.session.close()
